📝 Description

The memory bank models currently keep training embeddings, model's stored embeddings on device instead of releasing.

This can be seen most dramatically in Patchcore and PaDiM. We expect peak memory to drop in validation but actually we keep the training embeddings in memory:

As an example here is Patchcore's lightning fit function, and I've commented where we should delete duplicate embeddings:

def fit(self) -> None:
        """Apply subsampling to the embedding collected from the training set.

        This method:
        1. Aggregates embeddings from all training batches
        2. Applies coreset subsampling to reduce memory requirements
        """
        logger.info("Aggregating the embedding extracted from the training set.")
        embeddings = torch.vstack(self.embeddings)

        logger.info("Applying core-set subsampling to get the embedding.")
        self.model.subsample_embedding(embeddings, self.coreset_sampling_ratio)

        # del self.embeddings <-- we should remove the extra ones

In addition to duplicate embeddings we should also be storing embeddings on the model not in lightning. So that the model can be used stand alone.

Lastly, the memory bank models are not designed for DDP or multi-gpu use. I have put default args to fix this.

• 🛠️ Fixes # (issue number)
• [Bug]: When using two GPUs to train Patchcore, only half of the memory_bank saved in the model checkpoint file is saved #2673
• GPU out of memory using PatchCore e PaDim #2752

✨ Changes

Select what type of change your PR is:

• 🚀 New feature (non-breaking change which adds functionality)
• 🐞 Bug fix (non-breaking change which fixes an issue)
• 🔄 Refactor (non-breaking change which refactors the code base)
• ⚡ Performance improvements
• 🎨 Style changes (code style/formatting)
• 🧪 Tests (adding/modifying tests)
• 📚 Documentation update
• 📦 Build system changes
• 🚧 CI/CD configuration
• 🔧 Chore (general maintenance)
• 🔒 Security update
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)

✅ Checklist

Before you submit your pull request, please make sure you have completed the following steps:

• 📚 I have made the necessary updates to the documentation (if applicable).
• 🧪 I have written tests that support my changes and prove that my fix is effective or my feature works (if applicable).
• 🏷️ My PR title follows conventional commit format.

For more information about code review checklists, see the Code Review Checklist.